#include <Dynamics/Fem/fem_cuda.h>
#include <Utilities/glm_ext.h>
#include <Utilities/cuda_utilities.h>
#include <Dynamics/Fem/shape_function.h>
#include <cuda_runtime.h>
#include <cuda_block_size_set.h>
using namespace std;

namespace PhysLeo {

template<typename T>
FemCuda<T>::FemCuda():FemCpu<T>()
{
}

template<typename T>
void FemCuda<T>::allocateResource(SimData<T>& simData)
{
    //the same as FemCpu do, just read and set simulation data. and then compute the shapematrix
    FemCpu<T>::allocateResource(simData);
}

template<typename T>
__global__ static void K_FemCuda_cubicShapeMatrix(int* ele_ptr, glm::tvec3<T>* rest_pos_ptr, glm::tmat8x3<T>* mat_ptr,const int quadrature_points_, const int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    const auto ele_id = pid;

    //gaussian quadrature points integration information
    int sample_num[2] = { 1,8 };
    auto a = static_cast<T>(0.57735);
    glm::tvec3<T> one_points_pos[1] = { glm::tvec3<T>(0,0,0) };
    glm::tvec3<T> two_points_pos[8] = { glm::tvec3<T>(-a,-a,-a), glm::tvec3<T>(-a,-a,a), glm::tvec3<T>(-a,a,-a),  glm::tvec3<T>(-a,a,a),
        glm::tvec3<T>(a,-a,-a), glm::tvec3<T>(a,-a,a), glm::tvec3<T>(a,a,-a),glm::tvec3<T>(a,a,a) };
    glm::tvec3<T>* sample_point[2] = { one_points_pos, two_points_pos };

    glm::tmat8x3<T> rest_shape;
    for (auto i = 0; i < 8; ++i) {
        auto vert_index = ele_ptr[ele_id * 8 + i];
        rest_shape[i] = rest_pos_ptr[vert_index];
    }

    for (auto i = 0; i < sample_num[quadrature_points_ - 1]; ++i)
    {
        glm::tmat8x3<T> s;
        cubicShapeDerivatives(sample_point[quadrature_points_ - 1][i], s);
        mat_ptr[ele_id*sample_num[quadrature_points_ - 1] + i] = s * glm::inverse(rest_shape*s);
    }
}

template<typename T>
void FemCuda<T>::cubicShapeMatrix()
{
    int sample_num[2] = { 1,8 };

    int ele_num = FemCpu<T>::ptr_ele_volume_->size();
    FemCpu<T>::ptr_shape_matrix_ = make_shared<BufferData<T>>(3 * 8 * ele_num*sample_num[FemCpu<T>::quadrature_points_]);
    FemCpu<T>::ptr_shape_matrix_->allocateGpu();

    auto mat_ptr = reinterpret_cast<glm::tmat8x3<T>*>(FemCpu<T>::ptr_shape_matrix_->ptrGpu());
    auto ele_ptr = FemCpu<T>::ptr_ele_->ptrGpu();
    auto rest_pos_ptr = FemCpu<T>::ptr_rest_pos_->ptrGpu();

    K_FemCuda_cubicShapeMatrix<T> << <numBlocks(ele_num), BLOCK_SIZE >> > (ele_ptr, rest_pos_ptr, mat_ptr, FemCpu<T>::quadrature_points_, ele_num);
    synchronCheck;
}

template<typename T>
__global__ void static K_FemCuda_tetShapeMatrix(int* ele_ptr, glm::tvec3<T>* rest_pos_ptr, glm::tmat3x4<T>* mat_ptr, const int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    auto ele_id = pid;

    glm::tmat4x3<T> rest_shape;  //3行4列矩阵
    for (auto i = 0; i < 4; ++i) {
        auto vert_index = ele_ptr[4 * ele_id + i];
        rest_shape[i] = rest_pos_ptr[vert_index];
    }

    //实际上是 4行3列的矩阵
    glm::tmat3x4<T> s(-1.0, 1.0, 0.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0, 0.0, 0.0, 1.0);
    mat_ptr[ele_id] = s * glm::inverse(rest_shape * s);
}

template<typename T>
void FemCuda<T>::tetShapeMatrix()
{
    int ele_num = FemCpu<T>::ptr_ele_volume_->size();
    FemCpu<T>::ptr_shape_matrix_ = make_shared<BufferData<T>>(3 * 4 * ele_num);
    FemCpu<T>::ptr_shape_matrix_->allocateGpu();
    
    auto mat_ptr = reinterpret_cast<glm::tmat3x4<T>*>(FemCpu<T>::ptr_shape_matrix_->ptrGpu());
    auto ele_ptr = FemCpu<T>::ptr_ele_->ptrGpu();
    auto rest_pos_ptr = FemCpu<T>::ptr_rest_pos_->ptrGpu();

    K_FemCuda_tetShapeMatrix<T> << <numBlocks(ele_num), BLOCK_SIZE >> > (ele_ptr, rest_pos_ptr, mat_ptr, ele_num);
    synchronCheck;
}

template<typename T>
__global__ static void K_FemCuda_cubicInternalForce(int* ele_ptr, glm::tvec3<T>* pos_ptr, glm::tvec3<T>* force_ptr, glm::tmat8x3<T>* mat_ptr, int quadrature_points, T* ele_volume_ptr, T* ele_material_ptr, FemMaterialType fem_material_type, bool homogeneous, const int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    auto ele_id = pid;

    int sample_num[2] = { 1,8 };
    T one_point_weight[1] = { 1.0 };
    T two_point_weight[8] = { 0.125,0.125,0.125,0.125, 0.125,0.125, 0.125,0.125 };
    T* sample_weight[2] = { one_point_weight,two_point_weight };

    //material
    const int num_material_parameter = FemMaterial<T>::numMaterialParams(fem_material_type);
    ele_material_ptr = ele_material_ptr + (homogeneous ? 0 : ele_id * num_material_parameter);

    glm::tmat8x3<T> current_shape;
    for (auto i = 0; i < 8; ++i)
    {
        auto vert_index = ele_ptr[ele_id * 8 + i];
        current_shape[i] = pos_ptr[vert_index];
    }

    //force
    glm::tmat8x3<T> total_force(0.0);
    for (auto i = 0; i<sample_num[quadrature_points - 1]; ++i)
    {
        glm::tmat8x3<T> rest_shape = mat_ptr[ele_id*sample_num[quadrature_points - 1] + i];
        glm::tmat3x3<T> deform_gradient = current_shape * rest_shape;

        //solve invertible problem
        if (glm::determinant(deform_gradient) < 0)
        {
            glm::tvec3<T> col = deform_gradient[0];
            deform_gradient[0] = deform_gradient[1];
            deform_gradient[1] = col;
        }

        glm::tmat3x3<T> stress = FemMaterial<T>::firstPiolaStress(deform_gradient, ele_material_ptr, fem_material_type);
        total_force = total_force + static_cast<T>(-1) * sample_weight[quadrature_points - 1][i] * ele_volume_ptr[ele_id] * stress * rest_shape;
    }

    //atomic add force
    for (auto i = 0; i < 8; ++i) {
        auto vert_index = ele_ptr[ele_id * 8 + i];
        atomicAdd(&(force_ptr[vert_index][0]), total_force[i][0]);
        atomicAdd(&(force_ptr[vert_index][1]), total_force[i][1]);
        atomicAdd(&(force_ptr[vert_index][2]), total_force[i][2]);
    }
}

template<typename T>
void FemCuda<T>::cubicInternalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    int ele_num = FemCpu<T>::ptr_ele_volume_->size();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto force_ptr = ptr_force->ptrGpu();
    auto mat_ptr = reinterpret_cast<glm::tmat8x3<T>*>(FemCpu<T>::ptr_shape_matrix_->ptrGpu());
    auto ele_ptr = FemCpu<T>::ptr_ele_->ptrGpu();
    auto ele_volume_ptr = FemCpu<T>::ptr_ele_volume_->ptrGpu();
    auto ele_material_ptr = FemCpu<T>::ptr_ele_material_->ptrGpu();

    K_FemCuda_cubicInternalForce<T> << <numBlocks(ele_num), BLOCK_SIZE >> > (ele_ptr, pos_ptr, force_ptr, mat_ptr, FemCpu<T>::quadrature_points_, ele_volume_ptr, ele_material_ptr, FemCpu<T>::fem_material_type_, FemCpu<T>::homogeneous_, ele_num);
    synchronCheck;
}

template<typename T>
__global__ static void K_FemCuda_tetInternalForce(int* ele_ptr, glm::tvec3<T>* pos_ptr, glm::tvec3<T>* force_ptr, glm::tmat3x4<T>* mat_ptr, T* ele_volume_ptr, T* ele_material_ptr, FemMaterialType fem_material_type, bool homogeneous, const int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    auto ele_id = pid;

    //material
    const int num_material_parameter = FemMaterial<T>::numMaterialParams(fem_material_type);
    ele_material_ptr = ele_material_ptr + (homogeneous ? 0 : ele_id * num_material_parameter);

    glm::tmat4x3<T> current_shape;  //3行4列矩阵
    for (auto i = 0; i < 4; ++i) {
        auto vert_index = ele_ptr[ele_id * 4 + i];
        current_shape[i] = pos_ptr[vert_index];
    }

    glm::tmat3x4<T> rest_shape = mat_ptr[ele_id];
    glm::tmat3x3<T> deform_gradient = current_shape * rest_shape;

    //solve invertible problem
    if (glm::determinant(deform_gradient) < 0)
    {
        glm::tvec3<T> col = deform_gradient[0];
        deform_gradient[0] = deform_gradient[1];
        deform_gradient[1] = col;
    }

    glm::tmat3x3<T> stress = FemMaterial<T>::firstPiolaStress(deform_gradient, ele_material_ptr, fem_material_type);
    glm::tmat4x3<T> total_force = -ele_volume_ptr[ele_id] * stress * glm::transpose(rest_shape);

    //add force
    for (auto i = 0; i < 4; ++i) {
        auto vert_index = ele_ptr[ele_id * 4 + i];
        atomicAdd(&(force_ptr[vert_index][0]), total_force[i][0]);
        atomicAdd(&(force_ptr[vert_index][1]), total_force[i][1]);
        atomicAdd(&(force_ptr[vert_index][2]), total_force[i][2]);
    }
}

template<typename T>
void FemCuda<T>::tetInternalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    int ele_num = FemCpu<T>::ptr_ele_volume_->size();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto force_ptr = ptr_force->ptrGpu();
    auto mat_ptr = reinterpret_cast<glm::tmat3x4<T>*>(FemCpu<T>::ptr_shape_matrix_->ptrGpu());
    auto ele_ptr = FemCpu<T>::ptr_ele_->ptrGpu();
    auto ele_volume_ptr = FemCpu<T>::ptr_ele_volume_->ptrGpu();
    auto ele_material_ptr = FemCpu<T>::ptr_ele_material_->ptrGpu();

    K_FemCuda_tetInternalForce<T> << <numBlocks(ele_num), BLOCK_SIZE >> > (ele_ptr, pos_ptr, force_ptr, mat_ptr, ele_volume_ptr, ele_material_ptr, FemCpu<T>::fem_material_type_, FemCpu<T>::homogeneous_, ele_num);
    synchronCheck;
}

template class FemCuda<float>;
template class FemCuda<double>;

}